#!/usr/bin/env python3
import queue
import numpy as np
import rclpy
from rclpy.node import Node
from tf2_ros import TransformBroadcaster, TransformListener, Buffer, TransformStamped
from sophuspy import SE3
from transforms3d import quaternions as qua
from transforms3d import euler

"""
用于自动标定二维码位置到车体坐标系统的转移矩阵
坐标系统应该具有如下形式：
backCar:chassis
    -> con
    -> camera_mount ->camera -> tag
forwardCar:chassis
    -> rec
    -> stag [wait to calibrate]
tag<-==->stag
con<-==->rec
操作流程：
1 完成两车刚性连接
2 发送两车urdf坐标tf
3 启动tag检测节点

TODO average tag pose
"""


def preprocess_prefix(prefix: str) -> str:
    if len(prefix) == 0:
        return prefix
    if prefix[0] == "/":
        prefix = prefix[1:]
    if prefix[-1] != "/":
        prefix = prefix + "/"
    return prefix


def GetSE3(transform_stamped: TransformStamped) -> SE3:
    t = np.array(
        [
            transform_stamped.transform.translation.x,
            transform_stamped.transform.translation.y,
            transform_stamped.transform.translation.z,
        ]
    )
    q = np.array(
        [
            transform_stamped.transform.rotation.w,
            transform_stamped.transform.rotation.x,
            transform_stamped.transform.rotation.y,
            transform_stamped.transform.rotation.z,
        ]
    )
    return SE3(qua.quat2mat(q), t)


def se3_to_string(tf: SE3) -> str:
    return f"\n translation:{tf.translation()} \n rotation:{euler.mat2euler(tf.rotationMatrix())}"


class AutoCalibrate(Node):
    def __init__(self):
        super().__init__("auto_calibrate")
        self.get_logger().info("auto_calibrate node started")
        self.declare_parameter("forward_prefix", "robot2")
        self.declare_parameter("back_prefix", "robot1")
        self.declare_parameter("forward_part", "con_1")
        self.declare_parameter("back_part", "rec_1")
        self.forward_prefix: str = self.get_parameter("forward_prefix").value
        self.back_prefix: str = self.get_parameter("back_prefix").value
        self.forward_prefix = preprocess_prefix(self.forward_prefix)
        self.back_prefix = preprocess_prefix(self.back_prefix)

        self.forward_part = self.get_parameter("forward_part").value
        self.back_part = self.get_parameter("back_part").value

        self.declare_parameter("sample", 10)
        self.sample: int = self.get_parameter("sample").value

        self.tf_broadcaster = TransformBroadcaster(self)
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)
        self.buffer = queue.Queue(self.sample)

        self.create_timer(0.5, self.on_calibrate)

    def on_calibrate(self):
        tf = self.calibrate_tf()
        if tf is None:
            return
        #if self.buffer.qsize() >= self.sample:
        #   self.buffer.get()
        #self.buffer.put(tf)
        #tf = self.get_avr_tf()
        self.get_logger().info("calibrate:" + se3_to_string(tf))

    def calibrate_tf(self) -> SE3:
        try:
            chassis_back_to_camera = GetSE3(
                self.tf_buffer.lookup_transform(
                    self.B("chassis"), self.B("camera"), rclpy.time.Time()
                )
            )
            camera_to_tag = GetSE3(
                self.tf_buffer.lookup_transform(
                    self.B("camera"), self.B("tag"), rclpy.time.Time()
                )
            )
            back_chassis_to_con = GetSE3(
                self.tf_buffer.lookup_transform(
                    self.B("chassis"), self.B(self.back_part), rclpy.time.Time()
                )
            )
            rec_to_forward_chassis = GetSE3(
                self.tf_buffer.lookup_transform(
                    self.F(self.forward_part), self.F("chassis"), rclpy.time.Time()
                )
            )
        except Exception as e:
            self.get_logger().warn(f"Failed to lookup transform: {e}")
            return None
        back_chassis_to_tag = chassis_back_to_camera * camera_to_tag
        back_chassis_to_forward_chassis = back_chassis_to_con * rec_to_forward_chassis
        self.get_logger().info(
            f"back_chassis_to_forward_chassis: {se3_to_string(back_chassis_to_forward_chassis)}"
        )
        return  back_chassis_to_forward_chassis*back_chassis_to_tag

    def F(self, tf_name: str) -> str:
        return self.forward_prefix + tf_name

    def B(self, tf_name: str) -> str:
        return self.back_prefix + tf_name


def main(args=None):
    rclpy.init(args=args)
    node = AutoCalibrate()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()


if __name__ == "__main__":
    main()
